#!/usr/env python
# Rob Dennis - Arclite EMP
# TVDB API
# Project Homepage: http://code.google.com/p/keasby/

#--- Python imports
import collections
import logging
import os
import re
import time
import urllib2
import unittest
import socket
import sys
import zipfile

from contextlib import closing
from os.path import join as pjoin
from random import choice
from urllib import pathname2url, urlencode, urlretrieve
from xml.etree import ElementTree as ET

#--- Third Party imports
from elixir import *
from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound

log = logging.getLogger("tvdb.api")

socket.setdefaulttimeout(10)

class TvDbApi(object):
    """
    This class is written for interacting with thetvdb.com's API
    """

    base_url = "http://thetvdb.com"
    
    def_order = "TVDB API Default Order"
    dvd_order = "TVDB API DVD Order"
    abs_order = "TVDB Absolute Order"
    
    orders = (def_order, dvd_order, abs_order)

    class GetOneMixin(object):
        @classmethod
        def get_one(cls, **kwargs):
            """
            similar to Entity's built-in get and get_by functions\
            added for additional options when used
            """
            try:
                return cls.query.filter_by(**kwargs).one()
            except NoResultFound, msg:
                _error_string = "No %s query found using: %s"%(cls.__name__,
                                                               str(kwargs))
                raise TvDbApiGetOneQueryEmptyError(_error_string)
            except MultipleResultsFound, msg:
                _error_string = "".join(["%s query found >1"%cls.__name__,
                                         "results with %s:"%str(kwargs),
                                     str(cls.query.filter_by(**kwargs).all())])
                raise TvDbApiGetOneQueryMultipleError(_error_string)


    class TvDbApiInfo(Entity, GetOneMixin):
        update_time = Field(Float, default=0.0)

        series = OneToMany("Series")
        map = OneToMany("SearchMap")
       
        def __repr__(self):
            return "<TvDbApiInfo>"


    class Series(Entity, GetOneMixin):
        id = Field(Integer, primary_key=True)
        Actors = Field(Text)
        Airs_DayOfWeek = Field(Text)
        Airs_Time = Field(Text)
        ContentRating = Field(Text)
        FirstAired = Field(Text)
        Genre = Field(Text)
        IMDB_ID = Field(Text)
        Language = Field(Text)
        Network = Field(Text)
        NetworkID = Field(Text)
        Overview = Field(Text)
        Rating = Field(Float)
        Runtime = Field(Text)
        SeriesID = Field(Integer)
        SeriesName = Field(Text)
        Status = Field(Text)
        added = Field(Text)
        addedBy = Field(Text)
        banner = Field(Text)
        fanart = Field(Text)
        lastupdated = Field(Float, default=0.0)
        poster = Field(Text)
        zap2it_id = Field(Text)
        #api-specific
        will_update = Field(Boolean, default=False)

        episodes = OneToMany("Episode")
        tvdb = ManyToOne("TvDbApiInfo")

        field_names = ("id",
                       "Actors",
                       "Airs_DayOfWeek",
                       "Airs_Time",
                       "ContentRating",
                       "FirstAired",
                       "Genre",
                       "IMDB_ID",
                       "Language",
                       "Network",
                       "NetworkID",
                       "Overview",
                       "Rating",
                       "SeriesID",
                       "SeriesName",
                       "Status",
                       "added",
                       "addedBy",
                       "banner",
                       "fanart",
                       "lastupdated",
                       "poster",
                       "zap2it_id")

        def __repr__(self):
            return "<Series %d: %s>"%(self.id, self.SeriesName)


    class Episode(Entity, GetOneMixin):
        id = Field(Integer, primary_key=True)
        Combined_episodenumber = Field(Integer)
        Combined_season = Field(Integer)
        DVD_chapter = Field(Integer)
        DVD_discid = Field(Integer)
        DVD_episodenumber = Field(Integer)
        DVD_season = Field(Integer)
        Director = Field(Text)
        EpImgFlag = Field(Text)
        EpisodeName = Field(Text)
        EpisodeNumber = Field(Integer)
        FirstAired = Field(Text)
        GuestStars = Field(Text)
        IMDB_ID = Field(Text)
        Language = Field(Text)
        Overview = Field(Text)
        ProductionCode = Field(Text)
        Rating = Field(Float)
        SeasonNumber = Field(Integer)
        Writer = Field(Text)
        absolute_number = Field(Integer)
        filename = Field(Text)
        lastupdated = Field(Float, default=0.0)
        seasonid = Field(Integer)
        seriesid = Field(Integer)

        series = ManyToOne("Series")

        field_names = ("id",
                       "Combined_episodenumber",
                       "Combined_season",
                       "DVD_chapter",
                       "DVD_discid",
                       "DVD_episodenumber",
                       "DVD_season",
                       "Director",
                       "EpImgFlag",
                       "EpisodeName",
                       "EpisodeNumber",
                       "FirstAired",
                       "GuestStars",
                       "IMDB_ID",
                       "Language",
                       "Overview",
                       "ProductionCode",
                       "Rating",
                       "SeasonNumber",
                       "Writer",
                       "absolute_number",
                       "filename",
                       "lastupdated")

        def __repr__(self):
            return "<Episode %d: %s - S%dE%d - %s>"%(self.id,
                                                  self.series.SeriesName,
                                                  self.SeasonNumber,
                                                  self.EpisodeNumber,
                                                  self.EpisodeName)


    class SearchMap(Entity, GetOneMixin):
        name = Field(Text, primary_key=True)
        id = Field(Integer)

        tvdb = ManyToOne("TvDbApiInfo")
        
    
    class QuerySet(object):
        """
        Holds the results of a previous query for display of results\
        and chaining as in bracket/getitem implementations
        """
        
        def __init__(self, api,
                    order=None,
                    series=None):
            
            self.api = api
            self.order = order
            self.series = series
            
            assert self.order in TvDbApi.orders, "unknown order %s" % self.order

            if order == TvDbApi.dvd_order or order == TvDbApi.def_order:
                self.season, self.episode = (None, None)
            else:
                self.episode = None
            
        def __getitem__(self, argument):
            
            if not self.series:
                self.series = argument
            elif not hasattr(self, "season") and not self.episode:
                #no season used in ordering type, eg Absolute Ordering
                self.episode = argument
                return self._find()
            elif not self.season:
                #season not yet set
                self.season = argument
            elif self.season and not self.episode:
                #already set season and providing episode, the natural end-state
                self.episode = argument
                return self._find()
            else:
                raise TvDbApiError("unknown state"+
                                    " in getitem: %s"%str(self.__dict__))
                
            return self
            
        @property
        def numbers(self):
            """
            returns numbers value conforming to type expected by TvDbApi.find
            """
            
            if (self.order == TvDbApi.dvd_order or 
                self.order == TvDbApi.def_order):

                return (self.season, self.episode)
            else:

                return self.episode            
         
        def _find(self):
            """
            used as a method of returning usable information and stop\
            chaining QuerySets
            """
            return self.api.find(self.series, self.numbers, self.order)

        def __iter__(self):
        
            return self._find()
            
        def __repr__(self):
        
            attr_list = ["%s=%s" % (attr, getattr(self, attr)) 
                            for attr in ("order", "series", "season", "episode")
                            if hasattr(self, attr) and getattr(self, attr)]

            return "<QuerySet: %s>" % ", ".join(attr_list)
                
            
            
            

    def __init__(self, api_key, bind=None,
                       cache=pjoin("cache", "tvdb")):
        """
        Instantiate's class attributes for caching, interaction with data,\
        and for the initialization of mirrors and updates since last\
        interaction. Setting bind to None will default to placing the\
        cache directory
        """

        self.api_key = api_key

        # init mirrors
        self.mirrors = dict(xml=[self.base_url],
                            banner=[self.base_url],
                            zip=[self.base_url])

        # init language
        self.language = u"en"


        # init update time
        self.update_time = time.gmtime(0)

        self.bind = bind
        self.cache = cache

        self.setup()

#---URL Interaction
#--Queries

    def query_GetSeries(self, seriesname, language=None):
        """
        http://thetvdb.com/wiki/index.php/API:GetSeries
        """

        if language is None:
            data = urlencode(dict(seriesname=seriesname), doseq=True)
        else:
            data = urlencode(dict(language=language,
                                    seriesname=seriesname),
                                    doseq=True)

        url = TvDbApi._make_url(self.base_url, "api",
                                "GetSeries.php?%s"%data)

        return url

    def query_Updates(self, time=None, type=None):
        """
        http://thetvdb.com/wiki/index.php/API:Updates\
        type must be one of: [none|series|episode|all]\
        time is: epoch time since last update (limited to 1000 results\
                 or 30 days). Optional only if type is "none"
        """

        if isinstance(type, basestring):
            if type.lower() not in ("none", "series", "all", "episode"):
                raise TvDbApiError("type must be one of: "+
                                    "[none|series|episode|all]")

        if all([time, type]):
            data = urlencode(dict(time=time, type=type), doseq=True)
        elif time is not None:
            data = urlencode(dict(time=time))
        elif type is not None:
            data = urlencode(dict(type=type))
        else:
            raise TvDbApiError("time and type can't both be None")

        return TvDbApi._make_url(self.base_url, "api", "Updates.php?%s"%data)

#---Public-facing

    def add_series(self, *seriesids):
        """
        Adds a seriesid to the set to download from thetvdb.com
        """

        for seriesid in seriesids:
            assert isinstance(seriesid, int), ("can't add id of type" + 
                                                "%s"%type(seriesid))
            if seriesid not in self.seriesid_list:
                self.info.series.append(TvDbApi.Series(id=seriesid,
                                                        will_update=True))

        session.commit()

    def update(self, *seriesids):
        """
        downloads the information necessary to keep a series updated
        """

        self.add_series(*seriesids)

        if seriesids:
            update_list = seriesids
        else:
            update_list = self.seriesid_list

        for seriesid in update_list:
            #select correct series

            series = TvDbApi.Series.query.filter_by(id=seriesid).one()
            log.debug("selected series %s", series)

            if series.SeriesName is None:
                needs_full_update = True
            else:
                log.debug("Does not need full update")
                needs_full_update = False

            if needs_full_update:
                log.debug("full update required")

                path = self._make_fsr_zip_path(seriesid)
                url = self._make_api_url(self.get_mirror("zip"),
                                         pathname2url(path))

                path, _ = self._urlretrieve(url, path)

                try:
                    series_zip = zipfile.ZipFile(path, "r")
                except zipfile.BadZipFileError:
                    log.error("%s not a zipfile, dumping text:\n%s",
                                path, open(path, "r").read())
                except Exception:
                    log.error("error opening zipfile %s", path)
                    raise

                with closing(series_zip.open("%s.xml"%self.language)
                             ) as zipped_xml:

                    tree = ET.parse(zipped_xml)

                    series.from_dict(self._make_node_dict(series.field_names,
                                                          tree.find("Series")))

                    for episode_node in tree.getiterator("Episode"):
                        names = TvDbApi.Episode.field_names
                        episode_dict = self._make_node_dict(names, episode_node)

                        series.episodes.append(TvDbApi.Episode(**episode_dict))

        session.commit()

    def search(self, seriesname):
        """
        search tvdb for series IDs corresponding to a provided seriesname and
        return the list of series IDs found
        """

        url = self.query_GetSeries(seriesname, self.language)

        path = pjoin(self.cache, "query_GetSeries-%s.xml"%seriesname)

        path, _ = self._urlretrieve(url, path)

        tree = ET.parse(path)

        seriesid_list = []

        for series_node in tree.getiterator("Series"):
            series_dict = self._make_node_dict(TvDbApi.Series.field_names,
                                               series_node)

            seriesid_list.append(int(series_dict["id"]))


        self._select_seriesname_map(seriesname, seriesid_list)

        return seriesid_list
    
    def find(self, series, numbers, order=None):
        """
        return the episode based on the supplied:\
        SeriesName or seriesid\
        and, depending on the ordering method specified:\
        season number (int) and episode number (int) in broadcast order, or\
        season number (int) and episode number (int) in DVD order, or\
        absolute number (int).\
        finally, the last, optional argument is one of:\
        TvDbApi.def_order, TvDbApi.dvd_order, or TvDbApi.abs_order\
        if not specified, use TvDbApi.def_order\
        assumes:\
        series is a string (SeriesName) or int (seriesid)\
        numbers is a length-two iterable of ints (DVD or Def order) or an\
        int (Absolute order)
        """
        
        absolute_number = None
        
        if not isinstance(numbers, tuple):
            #in order to have a broad type signature, numbers could\
            #be a tuple (def/dvd order) or int (abs order)
            absolute_number = numbers
        
        if order is None:
            order = TvDbApi.def_order
        assert order in TvDbApi.orders, "%s is unknown order type"%order
        
        if isinstance(numbers, int):         
            if order != TvDbApi.abs_order:
                log.warning("assuming absolute order due to only 1"+
                            " number supplied: %s-%d", series, absolute_number)
                order = TvDbApi.abs_order

        elif len(numbers) != 2:
            raise TvDbApiError("numbers must be of length 2: %s"%str(numbers))
            
        log.debug("using ordering: %s", str(order))
        
        if isinstance(series, int):
            log.debug("searching using seriesid %d", series)

            series_query = dict(id=series)
            
        else:
            log.debug("searching using SeriesName %s", series)
            series_query = dict(SeriesName=series)   
            
        found_series = self._find_series(series_query)
        
        if not order == TvDbApi.abs_order and not any(numbers):
            #unable to narrow down from series, return it
            return found_series
            
        elif order == TvDbApi.abs_order and not absolute_number:
            #unable to narrow down from series, return it
            return found_series
        
        if order is TvDbApi.dvd_order:
            return self._find_dvd_order(found_series, *numbers)
        elif order is TvDbApi.abs_order:
            return self._find_abs_order(found_series, absolute_number)
        else:
            return self._find_def_order(found_series, *numbers)
 
    def _find_series(self, series_query):
        """
        Returns the TvDbApi.Series instance for the given query
        """
        try:
            found_series = self._query_series(**series_query)
        except TvDbApiGetOneQueryEmptyError:
            log.info("nothing found using %s, updating", str(series_query))
            
            if "id" in series_query:
                self.update(series_query["id"])
            else:
                try:
                    name = series_query["SeriesName"]
                    id = TvDbApi.SearchMap.get_one(name=name).id
                    
                except TvDbApiGetOneQueryEmptyError:
                    log.info("Seriesname %s not yet queried", name)                    
                    self.search(name)
                    id = TvDbApi.SearchMap.get_one(name=name).id
                    self.update(id)
                    
                    
            log.debug("attempting to find series with updated data")
            return self._find_series(series_query)
                    
        return found_series
         
    def _find_def_order(self, series, season_number, episode_number=None):
        """
        Return the TvDbApi.Episode instance based on the given numbers
        """
        
        assert season_number
        
        if episode_number:
            return self._query_episode(series=series,
                                       SeasonNumber=season_number,
                                       EpisodeNumber=episode_number)

        else:
            return self._query_season(series=series,
                                      SeasonNumber=season_number) 
        
    def _find_dvd_order(self, series, season_number, episode_number=None):
        """
        Return the TvDbApi.Episode instance based on the given numbers
        """    
        
        assert season_number
        
        if episode_number:
            return self._query_episode(series=series,
                                       DVD_season=season_number,
                                       DVD_episodenumber=episode_number)

        else:
            return self._query_season(series=series,
                                      DVD_season=season_number)  
            
        
    def _find_abs_order(self, series, absolute_number):
        """
        Return the TvDbApi.Episode instance based on the given number
        """    
        
        assert absolute_number
        
        return self._query_episode(series=series, 
                                   absolute_number=absolute_number)
           
    @staticmethod
    def _query_series(**series_query):
        return TvDbApi.Series.get_one(**series_query)

    @staticmethod
    def _query_episode(**episode_query):
        return TvDbApi.Episode.get_one(**episode_query)
    
    @staticmethod
    def _query_season(**season_query):
        return TvDbApi.Episode.query.filter_by(**season_query)
    
    def _select_seriesname_map(self, seriesname, seriesid_list):
    
        if not seriesid_list:
            raise TvDbApiError("no seriesids provided for name %s"%seriesname)

        self.info.map.append(TvDbApi.SearchMap(name=seriesname,
                                               id=seriesid_list[0]))
                                               
        session.commit()
        
        return seriesid_list[0]

    def _make_fsr_zip_path(self, seriesid):
        """
        return a relative path to the full series record zip path
        """

        assert isinstance(seriesid, int), "seriesid must be an integer"

        relative_fsr_path = pjoin("series",
                                    str(seriesid),
                                    "all",
                                    "%s.zip"%self.language)

        return relative_fsr_path        
        
#---Properties

    @property
    def seriesid_list(self):

        return [series.id for series
                    in TvDbApi.Series.query.filter_by(will_update=True).all()]


#---Utility
    
    def __getitem__(self, argument):
        
        if argument not in (self.dvd_order, self.def_order, self.abs_order):
            log.debug("assuming default order")
            order = self.def_order
            series = argument
            log.debug("querying seriesname: %s", series)
            
            return self.QuerySet(self, order=order, series=series)

        else:
            order = argument
            log.debug("using: %s",argument)
            
            return self.QuerySet(self, order=order)        
        

    def get_mirror(self, mirror_type):
        """
        Return a mirror url of the appropriate type, errors if an unknown or \
        uninitialized type is specified.\
        mirror types: xml, zip, banner
        """

        if mirror_type not in self.mirrors:
            raise TvDbApiError("invalid mirror type specified: %s", mirror_type)

        return choice(self.mirrors[mirror_type])

    @staticmethod
    def _make_node_dict(fields, node):
            field_gen = ((field, node.findtext(field))
                            for field in fields
                            if node.findtext(field))

            return dict(field_gen)

    def _make_api_url(self, url_netloc=None, *url_path):
        """
        Convenience function for interacting accessing a URL relative to\
        the tvdb API key
        """

        if url_netloc is None:
            url_netloc = self.base_url

        return TvDbApi._make_url(url_netloc, "api", self.api_key, *url_path)

    @staticmethod
    def _make_url(url_netloc, *url_path):
        """
        Appends all parts of url_path to the given url_netloc using "/" as\
        a delimeter
        """

        url_sequence = [url_netloc]
        url_sequence.extend(url_path)

        return "/".join(url_sequence)

    @staticmethod
    def _urlopen(url, data=None):
        """
        Opens a given url
        """

        if data is None:
            log.debug("attempting to open url: %s", url)
            return urllib2.urlopen(url)

        else:
            log.debug("attempting to open url: %s, data: %s", url, data)
            return urllib2.urlopen(url, data)

    def _urlretrieve(self, url, filename=None, reporthook=None, data=None):
        """
        http://docs.python.org/library/urllib.html#urllib.urlretrieve\
        modifed so that a provided filename will have the cache directory
        prepended
        """

        if filename is not None:
            path = pjoin(self.cache, filename)
            parent, _ = os.path.split(path)
            if not os.path.exists(parent):
                os.makedirs(parent)

        try:
            path, info = urlretrieve(url, path, reporthook, data)
            return path, info
        except Exception, msg:

            log.warning("returning cache for %s url\n(%s:%s)"%(url,
                                                        msg.__class__.__name__,
                                                        str(msg)))
                                                        
            if os.path.exists(path):
                return path, None
            else:
                raise TvDbApiError("unable to retrieve cached information"+
                                    "from url: %s\nand path: %s\n"%(url, path)+
                                    "%s"%str(msg))

#--Housekeeping

    def setup(self):
        """
        Create the database connection and populate
        """

        if not os.path.exists(self.cache):
            os.makedirs(self.cache)

        if self.bind is None:
            self.bind = "sqlite:///%s"%pjoin(self.cache, "tvdb.sqlite")

        #elixir setup steps
        metadata.bind = self.bind
        setup_all()
        create_all()

        self.info = TvDbApi.TvDbApiInfo()


    def purge(self, delete=False):
        """
        Remove existing tables and optionally remove the file
        """

        session.close()
        drop_all()

        if delete:
            metadata.bind.dispose()
            import string
            filename = string.split(self.bind, "///")[-1]
            if os.path.exists(filename):
                os.remove(filename)


    def __repr__(self):
        """
        Representation
        """

        return "<%s>"%self.__class__.__name__

    def __del__(self):
        """
        clean up on closure
        """

        if session:
            session.close()


class TvDbApiError(Exception):
    """
    Generic TvDbApi Exception
    """
    pass
    

class TvDbApiQueryError(TvDbApiError):
    """
    Generic exception for TvDbApi queries
    """
    pass
    
    
class TvDbApiGetOneQueryEmptyError(TvDbApiQueryError):
    """
    Exception for no results returned from a get_one query
    """
    pass
    
    
class TvDbApiGetOneQueryMultipleError(TvDbApiQueryError):
    """
    Exception for no results returned from a get_one query
    """
    pass
    
    
class TvDbApiSeriesNotFoundError(TvDbApiError):
    """
    Exception for a given series not found
    """
    pass
__all__ = ["TvDbApi", 
            "TvDbApiError", 
            "TvDbApiSeriesNotFoundError"]

